clear all; close all; clc;

%% Parameters
nu = 1;              % Diffusion coefficient
L = 1;               % Domain length
h = 1/20;            % Spatial step size
x = 0:h:L;           % Spatial grid
Nx = length(x);    % Number of interior points

% Time parameters
r_values = [1/2, 1];   % r = nu*dt/h^2 values to test
k_values = r_values*h^2/nu; % Corresponding time steps

% Number of steps to plot
steps_to_plot = [1, 2, 10];

%% Initial condition (triangular function)
phi = @(x) (x>=9/20 & x<1/2).*(20*(x-9/20)) + ...
           (x>=1/2 & x<11/20).*(-20*(x-11/20));
u0 = phi(x(1:end))'; % Interior points only

%% Compute and plot solutions
figure('Position', [100, 100, 1200, 600]);
titles = {'After 1 step', 'After 2 steps', 'After 10 steps'};
methods = {'FTCS r=1/2', 'FTCS r=1'};

for col = 1:3
    steps = steps_to_plot(col);
    
    % FTCS r=1/2 (first row)
    subplot(2, 3, col);
    r = 1/2;
    k = r*h^2/nu;
    u = u0;
    for n = 1:steps
        u = ftcs(u, r, Nx);
    end
    plot(x(1:end), u, 'b-', 'LineWidth', 2);
    title(titles{col});
    ylim([0 0.6]);
    if col == 1
        ylabel(methods{1});
    end
    grid on;
    
    
    % FTCS r=1 (second row)
    subplot(2, 3, col+3);
    r = 1;
    k = r*h^2/nu;
    u = u0;
    for n = 1:steps
        u = ftcs(u, r, Nx);
    end
    plot(x(1:end), u, 'g-', 'LineWidth', 2);
    
    if col == 1
        ylabel(methods{2});
    end
    grid on;
end

sgtitle('FTCS for Heat Equation');

%% Function Definitions

function u_new = ftcs(u_old, r, Nx)
    % Backward Time Centered Space method
    alpha = r;
    main_diag = (1-2*alpha)*ones(Nx,1);
    off_diag = alpha*ones(Nx-1,1);
    A = diag(main_diag) + diag(off_diag,1) + diag(off_diag,-1);
    u_new = A*u_old;
end